#!/usr/bin/env python
'''
SlicedIPScheduler class for co-scheduling reservations
across multiple resources using time-slicing and an integer program.

Because time is discretized into time slices, this scheduler requires
information about how to generate the slices, so its signature has one
more argument than usual. 

Author: Sotiria Lampoudi (slampoud@gmail.com)
'''

import math
import numpy as np
from adaptive_scheduler.kernel.scheduler import Scheduler
from adaptive_scheduler.utils import OptimizationType

class PossibleStart(object):
    def __init__(self, resource, slice_starts, internal_start, airmass_coefficient):
        self.resource = resource
        self.first_slice_start = slice_starts[0]
        self.all_slice_starts = slice_starts
        self.internal_start = internal_start
        self.airmass_coefficient = airmass_coefficient

    def __lt__(self, other):
        return self.first_slice_start < other.first_slice_start

    def __eq__(self, other):
        return self.first_slice_start == self.first_slice_start

    def __gt__(self, other):
        return self.first_slice_start > self.first_slice_start


class SlicedIPScheduler_v2(Scheduler):

    def __init__(self, compound_reservation_list,
                 globally_possible_windows_dict,
                 contractual_obligation_list,
                 slice_size_seconds):
        super().__init__(compound_reservation_list,
                           globally_possible_windows_dict,
                           contractual_obligation_list)
        # time_slicing_dict is a dictionary that maps: 
        # resource-> [slice_alignment, slice_length]
        #         self.resource_list = resource_list
        self.slice_size_seconds = slice_size_seconds
        self.time_slicing_dict = {}
        # these are the structures we need for the linear programming solver
        self.Yik = []  # maps idx -> [resID, window idx, priority, resource]
        self.aikt = {}  # maps slice -> Yik idxs
        self.schedulerIDstring = 'slicedIPscheduler'
        self.hashes = set()

        for r in self.resource_list:
            self.time_slicing_dict[r] = [0, self.slice_size_seconds]

    def hash_slice(self, start, resource, slice_length):
        string = "resource_" + resource + "_start_" + repr(start) + "_length_" + repr(slice_length)
        exists = string in self.hashes
        self.hashes.add(string)
        return string, exists

    def unhash_slice(self, mystr):
        l = mystr.split("_")
        return [l[1], int(l[3]), int(l[5])]

    def build_data_structures(self):
        # first we need to build up the list of discretized slices that each
        # reservation can begin in. These are represented as attributes
        # that get attached to the reservation object. 
        # The new attributes are:
        # slices_dict
        # internal_starts_dict
        # and the dicts are keyed by resource.
        # the description of slices and internal starts is in intervals.py
        for r in self.reservation_list:
            r.Yik_entries = []
            r.possible_starts = []
            for resource in sorted(r.free_windows_dict.keys()):
                r.possible_starts.extend(self.get_slices(r.free_windows_dict[resource], resource, r.duration, r.request))
            # reorder PossibleStarts
            r.possible_starts.sort()
            # build aikt
            w_idx = 0
            for ps in r.possible_starts:
                Yik_idx = len(self.Yik)
                r.Yik_entries.append(Yik_idx)
                # set the initial warm start solution
                scheduled = 0
                if r.previous_solution_reservation and r.previous_solution_reservation.scheduled_start == ps.internal_start and r.previous_solution_reservation.scheduled_resource == ps.resource:
                    scheduled = 1
                # now w_idx is the index into r.possible_starts, which have
                # been reordered by time.
                if r.request and r.request.optimization_type == OptimizationType.AIRMASS:
                    # Apply the airmass coefficient into the priority
                    priority = r.priority + ps.airmass_coefficient
                else:
                    # Add the earlier window optimization priority factor to the effective priority
                    priority = r.priority + (0.1 / (w_idx + 1.0))
                self.Yik.append([r.resID, w_idx, priority, ps.resource, scheduled])
                w_idx += 1
                # build aikt
                for s in ps.all_slice_starts:
                    key, exists = self.hash_slice(s, ps.resource, self.time_slicing_dict[ps.resource][1])
                    #                        if key in self.aikt:
                    if exists:
                        self.aikt[key].append(Yik_idx)
                    else:
                        self.aikt[key] = [Yik_idx]

    def unpack_result(self, r):
        #        print(r.xf)
        idx = 0
        for value in r.xf:
            if value == 1:
                resID = self.Yik[idx][0]
                start_idx = self.Yik[idx][1]
                resource = self.Yik[idx][3]
                reservation = self.get_reservation_by_ID(resID)
                # use the internal_start for the start
                start = reservation.possible_starts[start_idx].internal_start
                # the quantum is the length of all the slices we've occupied
                quantum = reservation.possible_starts[start_idx].all_slice_starts[-1] + \
                          self.time_slicing_dict[resource][1] - reservation.possible_starts[start_idx].first_slice_start
                reservation.schedule(start, quantum, resource, self.schedulerIDstring)
                self.commit_reservation_to_schedule(reservation)
            idx += 1
        return self.schedule_dict

    def get_slices(self, intervals, resource, duration, request):
        ''' Creates two things:
        * slices: list of lists. Each inner list is a window. The first
        element is the initial slice, and each subsequent slice is also
        occupied. All slices are aligned with slice_alignment, and are
        slice_length long.
        * internal_starts: list of values, one per inner list of slices.
        Each internal_start can be either equal to the corresponding
        slices[0], in which case it's not internal, or > than it, being
        internal.
        Returns: a list of PossibleStart objects'''

        ps_list = []
        # Make sure the resource is available.  If it is not in the time slicing dict, it's not available
        if resource in self.time_slicing_dict:
            slice_alignment = self.time_slicing_dict[resource][0]
            slice_length = self.time_slicing_dict[resource][1]
            slices = []
            internal_starts = []
            for t in intervals.toDictList():
                if t['type'] == 'start':
                    if t['time'] <= slice_alignment:
                        start = slice_alignment
                        internal_start = slice_alignment
                    else:
                        # figure out start so it aligns with slice_alignment 
                        start = int(slice_alignment + math.floor(
                            float(t['time'] - slice_alignment) / float(slice_length)) * slice_length)
                        # use the actual start as an internal start (may or may not align w/ slice_alignment)
                        internal_start = t['time']
                elif t['type'] == 'end':
                    if t['time'] < slice_alignment:
                        continue
                    while t['time'] - start >= duration:
                        tmp = range(start, internal_start + duration, slice_length)
                        slices.append(tmp)
                        internal_starts.append(internal_start)
                        start += slice_length
                        internal_start = start

            # Get the airmass coefficients for the internal starts if the request should be optimized by airmass
            if request and request.optimization_type == OptimizationType.AIRMASS:
                airmasses_at_times = request.get_airmasses_within_kernel_windows(resource)
                # use numpy to interpolate airmass values for the internal_starts times
                interpolated_airmasses = np.interp(internal_starts, airmasses_at_times['times'], airmasses_at_times['airmasses'])
            else:
                interpolated_airmasses = np.zeros(len(internal_starts))
            # return slices, internal_starts
            ps_list = []
            idx = 0
            for w in slices:
                ps_list.append(PossibleStart(resource, w, internal_starts[idx], interpolated_airmasses[idx]))
                idx += 1

        return ps_list
